import numpy as np
import matplotlib.pyplot as plt
import math

def getSin(amp, freq, phase, sampleList):
    return amp * np.sin(-2 * math.pi * freq * sampleList + phase)

def getCos(amp, freq, phase, sampleList):
    return amp * np.cos(-2 * math.pi * freq * sampleList + phase)

def denoise(arr, thresh):
    mask = arr > thresh
    return mask * arr

srate = 3000
t = np.linspace(0, 1, srate)

s1 = getSin(amp=1.5, freq=30, phase=0, sampleList=t)
s2 = getCos(amp=3, freq=5, phase=0, sampleList=t)
s3 = getSin(amp=10, freq=100, phase=0, sampleList=t)
s4 = getCos(amp=20, freq=120, phase=0, sampleList=t)

mixed_signal = s1 + s2 + s3 + s4

fCoefs = np.fft.fft(mixed_signal, srate)

amp_list = 2 * np.abs(fCoefs / srate)

freqs = np.fft.fftfreq(len(amp_list), 1/srate)
amp_shifted = np.fft.fftshift(amp_list)
freq_shift = np.fft.fftshift(freqs)

fig, ax = plt.subplots(2, 2, figsize=(12, 8))

# 绘制原始未过滤数据的复合波形时域图
ax[0, 0].plot(t[:500], mixed_signal[:500])
ax[0, 0].grid()
ax[0, 0].set_yticks([-20, 0, 20])
ax[0, 0].set_xlim(0, 500/srate)
ax[0, 0].set_xticks(np.arange(0, 501/srate, 100/srate))
ax[0, 0].set_xticklabels(np.arange(0, 501, 100))
ax[0, 0].relim()
ax[0, 0].autoscale_view(scalex=False, scaley=True)

# 绘制原始振幅频谱
ax[1, 0].stem(freq_shift, amp_shifted)
ax[1, 0].set_xlim([-150, 150])
ax[1, 0].set_ylim([-0.5, 21])
ax[1, 0].grid()

amp_shifted[(freq_shift > 110) | (freq_shift < -110)] = 0
amp_shifted = denoise(amp_shifted, 1)

# 绘制处理后的振幅频谱
ax[1, 1].stem(freq_shift, amp_shifted)
ax[1, 1].set_xlim([-150, 150])
ax[1, 1].set_ylim([-0.5, 21])
ax[1, 1].grid()

# 逆傅里叶变换以获得双重过滤后的时域波形
filtered_signal = np.fft.ifft(np.fft.ifftshift(amp_shifted * np.fft.fftshift(fCoefs / amp_list)), srate)

# 绘制双重过滤后的混合波形时域图
ax[0, 1].plot(t[:500], filtered_signal.real[:500])
ax[0, 1].grid()
ax[0, 1].set_yticks(np.arange(-15, 16, 5))
ax[0, 1].set_xlim(0, 500/srate)
ax[0, 1].set_xticks(np.arange(0, 501/srate, 100/srate))
ax[0, 1].set_xticklabels(np.arange(0, 501, 100))
ax[0, 1].relim()
ax[0, 1].autoscale_view(scalex=False, scaley=True)

plt.tight_layout()
plt.show()